import numpy as np
import pandas as pd
from math import radians, sin, cos, sqrt, atan2


# 节点数据
nodes_data = [
    {"id": 76, "name": "朝天门", "type": "街道", "population": 30718, "area": 1.742, "lon": 106.5826, "lat": 29.56288},
    {"id": 64, "name": "解放碑", "type": "街道", "population": 52518, "area": 1.32, "lon": 106.5726, "lat": 29.56024},
    {"id": 63, "name": "南纪门", "type": "街道", "population": 36752, "area": 1.404, "lon": 106.5714, "lat": 29.55322},
    {"id": 50, "name": "七星岗", "type": "街道", "population": 49189, "area": 0.7086, "lon": 106.5611, "lat": 29.55791},
    {"id": 33, "name": "菜园坝", "type": "街道", "population": 24817, "area": 4.013, "lon": 106.5407, "lat": 29.54703},
    {"id": 31, "name": "两路口", "type": "街道", "population": 35877, "area": 2.112, "lon": 106.5385, "lat": 29.55247},
    {"id": 46, "name": "大溪沟", "type": "街道", "population": 58955, "area": 1.514, "lon": 106.5562, "lat": 29.56623},
    {"id": 34, "name": "上清寺", "type": "街道", "population": 46825, "area": 1.642, "lon": 106.5427, "lat": 29.56242},
    {"id": 6, "name": "石油路", "type": "街道", "population": 63691, "area": 3.168, "lon": 106.5, "lat": 29.54295},
    {"id": 19, "name": "大坪", "type": "街道", "population": 62344, "area": 2.261, "lon": 106.517, "lat": 29.54181},
    {"id": 9, "name": "化龙桥", "type": "街道", "population": 28366, "area": 3.35, "lon": 106.5073, "lat": 29.55295},
    {"id": 1, "name": "沙坪坝区", "type": "其他行政区", "population": 1490700, "area": 395.8, "lon": 106.4891,
     "lat": 29.55567},
    {"id": 11, "name": "九龙坡区", "type": "其他行政区", "population": 1540200, "area": 432, "lon": 106.5093,
     "lat": 29.531},
    {"id": 30, "name": "江北区1", "type": "其他行政区", "population": 947400, "area": 220.77, "lon": 106.5385,
     "lat": 29.56378},
    {"id": 38, "name": "江北区2", "type": "其他行政区", "population": 947400, "area": 220.77, "lon": 106.5462,
     "lat": 29.56932},
    {"id": 51, "name": "江北区3", "type": "其他行政区", "population": 947400, "area": 220.77, "lon": 106.5626,
     "lat": 29.5658},
    {"id": 48, "name": "南坪区", "type": "其他行政区", "population": 1210000, "area": 262.43, "lon": 106.5592,
     "lat": 29.55205}
]

# 渝中区整体数据
yuzhong_data = {
    "population": 490052,
    "area": 23.24,
    "passenger_volume": 211490000  # 年客运量
}


def haversine_distance(lon1, lat1, lon2, lat2):
    """计算两点间的球面距离（单位：公里）"""
    R = 6371  # 地球半径（公里）

    lon1, lat1, lon2, lat2 = map(radians, [lon1, lat1, lon2, lat2])
    dlon = lon2 - lon1
    dlat = lat2 - lat1

    a = sin(dlat / 2) ** 2 + cos(lat1) * cos(lat2) * sin(dlon / 2) ** 2
    c = 2 * atan2(sqrt(a), sqrt(1 - a))
    distance = R * c

    return distance


def calculate_od_matrix(nodes_data, yuzhong_data):
    """计算OD需求矩阵"""
    n_nodes = len(nodes_data)

    # 计算距离矩阵
    distances = np.zeros((n_nodes, n_nodes))
    for i, node1 in enumerate(nodes_data):
        for j, node2 in enumerate(nodes_data):
            if i != j:
                distances[i, j] = haversine_distance(
                    node1['lon'], node1['lat'],
                    node2['lon'], node2['lat']
                )

    # 计算小时客运量
    daily_volume = yuzhong_data['passenger_volume'] / 365
    peak_period_ratio = 0.8
    peak_hours = 16
    hourly_volume = (daily_volume * peak_period_ratio ) / peak_hours

    # 计算产生和吸引系数
    production_indices = {}
    attraction_indices = {}

    # 计算街道总人口
    total_street_population = sum(node['population'] for node in nodes_data if node['type'] == '街道')

    # 区内外交通量分配比例
    inter_district_ratio = 0.1  # 区间交通比例
    intra_district_ratio = 1 - inter_district_ratio  # 区内交通比例

    # 计算产生和吸引系数
    for node in nodes_data:
        if node['type'] == '街道':
            # 街道的产生力和吸引力计算
            pop_ratio = node['population'] / total_street_population

            # 产生力计算
            production_indices[node['id']] = pop_ratio * intra_district_ratio

            # 吸引力计算 - 考虑区位特征
            base_attraction = pop_ratio * intra_district_ratio
            if node['name'] in ['解放碑', '朝天门']:
                attraction_weight = 2.0
            elif node['name'] in ['大坪', '石油路']:
                attraction_weight = 1.8
            elif node['name'] in ['两路口', '七星岗']:
                attraction_weight = 1.5
            else:
                attraction_weight = 1.0

            attraction_indices[node['id']] = base_attraction * attraction_weight

        else:
            # 其他行政区计算
            district_pop_ratio = node['population'] / yuzhong_data['population']
            base_index = district_pop_ratio * inter_district_ratio

            if '江北区' in node['name']:
                base_index = base_index / 3

            production_indices[node['id']] = base_index
            attraction_indices[node['id']] = base_index

    # 归一化系数
    total_production = sum(production_indices.values())
    total_attraction = sum(attraction_indices.values())

    for node_id in production_indices:
        production_indices[node_id] /= total_production
        attraction_indices[node_id] /= total_attraction

    # 计算OD矩阵
    od_matrix = np.zeros((n_nodes, n_nodes))

    # 修改后的距离衰减系数
    decay_factor = 0.15  # 降低衰减系数，使距离的影响减小

    for i, node1 in enumerate(nodes_data):
        for j, node2 in enumerate(nodes_data):
            if i != j:
                # 计算基础流量
                flow = hourly_volume * production_indices[node1['id']] * attraction_indices[node2['id']]
                # 使用修改后的距离衰减函数
                distance_decay = np.exp(-decay_factor * distances[i, j])
                od_matrix[i, j] = flow * distance_decay

    # 对矩阵进行放大，使总量接近hourly_volume
    total_flow = np.sum(od_matrix)
    scaling_factor = hourly_volume / total_flow
    od_matrix *= scaling_factor

    # 创建DataFrame
    node_names = [node['id'] for node in nodes_data]
    od_df = pd.DataFrame(od_matrix, index=node_names, columns=node_names)

    return od_df


# 计算OD矩阵
od_matrix = calculate_od_matrix(nodes_data, yuzhong_data)

# 设置显示格式
pd.set_option('display.float_format', lambda x: '%.1f' % x)
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

# 输出结果
print("渝中区及周边区域OD需求矩阵 (单位: 人次/小时)")
print("=" * 100)
print(od_matrix)

# 保存到CSV文件
od_matrix.to_csv('od_matrix_output.csv')